package main

import (
	"io"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"

	"gopkg.in/ini.v1"
)

func main() {
	// 读取配置文件
	cfg, err := ini.Load("config.ini")
	if err != nil {
		log.Fatalf("无法加载配置文件: %v", err)
	}

	// 从配置文件中获取目标服务器地址和协议类型
	target := cfg.Section("server").Key("target").String()
	protocol := cfg.Section("server").Key("protocol").String()

	// 解析目标URL
	url, err := url.Parse(target)
	if err != nil {
		log.Fatalf("解析URL错误: %v", err)
	}

	if protocol == "http" {
		// HTTP代理
		http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
			proxy := httputil.NewSingleHostReverseProxy(url)
			proxy.Director = func(req *http.Request) {
				req.Header = r.Header
				req.Host = url.Host
				req.URL.Scheme = url.Scheme
				req.URL.Host = url.Host
				req.URL.Path = r.URL.Path
			}
			proxy.ErrorHandler = func(rw http.ResponseWriter, req *http.Request, err error) {
				http.Error(rw, "代理错误", http.StatusBadGateway)
				log.Printf("代理错误: %v", err)
			}
			proxy.ServeHTTP(w, r)
		})

		// 启动HTTP服务器
		log.Println("正在监听端口 8080 并转发到", target)
		log.Fatal(http.ListenAndServe(":8080", nil))

	} else if protocol == "tcp" {
		// TCP代理
		listener, err := net.Listen("tcp", ":8080")
		if err != nil {
			log.Fatalf("监听端口错误: %v", err)
		}

		log.Println("正在监听TCP端口 8080 并转发到", target)
		for {
			conn, err := listener.Accept()
			if err != nil {
				log.Printf("接受连接错误: %v", err)
				continue
			}
			go handleTCPConnection(conn, target)
		}
	} else {
		log.Fatalf("不支持的协议: %s", protocol)
	}
}

// 处理TCP连接
func handleTCPConnection(conn net.Conn, target string) {
	defer conn.Close()

	// 连接到目标服务器
	targetConn, err := net.Dial("tcp", target)
	if err != nil {
		log.Printf("连接目标服务器错误: %v", err)
		return
	}
	defer targetConn.Close()

	// 转发数据
	go func() {
		_, _ = io.Copy(targetConn, conn)
	}()
	_, _ = io.Copy(conn, targetConn)
}
